package com.xiaoyu.awakening.ai.index.service.impl;

import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.xiaoyu.awakening.ai.index.config.AwkThreadPoolConfig;
import com.xiaoyu.awakening.ai.index.dao.AwkIndexTaskDao;
import com.xiaoyu.awakening.ai.index.entity.AwkIndexPipelineContext;
import com.xiaoyu.awakening.ai.index.entity.AwkIndexTask;
import com.xiaoyu.awakening.ai.index.entity.AwkIndexTaskItem;
import com.xiaoyu.awakening.ai.index.entity.req.AwkIndexTaskReq;
import com.xiaoyu.awakening.ai.index.entity.res.AwkIndexTaskStatusDTO;
import com.xiaoyu.awakening.ai.index.service.IAwkIndexDocumentService;
import com.xiaoyu.awakening.ai.index.service.IAwkIndexTaskItemService;
import com.xiaoyu.awakening.ai.index.service.IAwkIndexTaskService;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.UUID;
import java.util.stream.Collectors;

/**
 * 索引任务
 *
 * @author wangxiaoyu
 * @date 2025/6/25 14:57
 * @since v1.0
 */
@Service
public class AwkIndexTaskServiceImpl extends ServiceImpl<AwkIndexTaskDao, AwkIndexTask> implements IAwkIndexTaskService {

    @Resource
    private IAwkIndexTaskItemService awkIndexTaskItemService;

    @Resource
    private IAwkIndexDocumentService indexDocumentService;

    @Override
    public AwkIndexTaskStatusDTO commitTask(AwkIndexTaskReq awkIndexTaskReq) {
        //持久化索引任务
        AwkIndexTask awkIndexTask = new AwkIndexTask();
        awkIndexTask.setTaskId(UUID.randomUUID().toString());
        awkIndexTask.setFileTotalNum(awkIndexTaskReq.getTotal());
        awkIndexTask.setTaskStatus(0);
        this.save(awkIndexTask);

        //构建索引任务项
        List<AwkIndexTaskItem> indexTaskItems = awkIndexTaskReq.getFileList().stream().map(awkIndexDocDto -> new AwkIndexTaskItem() {
            {
                setTaskId(awkIndexTask.getId());
                setFileId(awkIndexDocDto.getFileId());
                setFilePath(awkIndexDocDto.getDownloadUrl());
                setIndexStatus(0);
            }
        }).collect(Collectors.toList());
        awkIndexTaskItemService.saveBatch(indexTaskItems);

        //提交任务
        AwkThreadPoolConfig.AWK_THREAD_POOL.get(AwkThreadPoolConfig.INDEX_TASK_THREAD_POOL_NAME).execute(() -> {

            indexTaskItems.forEach(indexTaskItem -> {
                AwkIndexPipelineContext indexPipelineContext = AwkIndexPipelineContext.builder()
                        .taskId(indexTaskItem.getTaskId())
                        .taskItemId(indexTaskItem.getId())
                        .fileId(indexTaskItem.getFileId())
                        .downloadUrl(indexTaskItem.getFilePath())
                        .build();
                Boolean indexPipelineStatus = indexDocumentService.indexDocumentPipeline(indexPipelineContext);
                if (indexPipelineStatus) {
                    indexTaskItem.setIndexStatus(1);
                    awkIndexTaskItemService.updateById(indexTaskItem);
                } else {
                    indexTaskItem.setIndexStatus(2);
                    awkIndexTaskItemService.updateById(indexTaskItem);
                }
            });

            //更新任务状态
            awkIndexTask.setTaskStatus(1);
        });

        //返回任务ID
        AwkIndexTaskStatusDTO awkIndexTaskStatusDTO = new AwkIndexTaskStatusDTO();
        awkIndexTaskStatusDTO.setTaskId(awkIndexTask.getTaskId());
        awkIndexTaskStatusDTO.setStatus(awkIndexTask.getTaskStatus());
        awkIndexTaskStatusDTO.setProcessPercent(0);
        return awkIndexTaskStatusDTO;

    }

    @Override
    public AwkIndexTaskStatusDTO queryTaskStatus(String taskId) {
        AwkIndexTask awkIndexTask = this.query().eq("task_id", taskId).one();
        if (awkIndexTask != null) {
            AwkIndexTaskStatusDTO awkIndexTaskStatusDTO = new AwkIndexTaskStatusDTO();
            awkIndexTaskStatusDTO.setTaskId(awkIndexTask.getTaskId());
            awkIndexTaskStatusDTO.setStatus(awkIndexTask.getTaskStatus());
            Integer processPercent = (awkIndexTask.getSuccessNum() + awkIndexTask.getFailNum()) / awkIndexTask.getFileTotalNum() * 100;
            awkIndexTaskStatusDTO.setProcessPercent(processPercent);
            return awkIndexTaskStatusDTO;
        }
        return null;
    }

    @Override
    public AwkIndexTask queryTaskByTaskId(String taskId) {
        AwkIndexTask task = this.query().eq("task_id", taskId).one();
        if (task != null) {
            List<AwkIndexTaskItem> taskItems = awkIndexTaskItemService.query().eq("task_id", task.getId()).list();
            task.setTaskItems(taskItems);
            return task;
        }
        return null;
    }


    @Override
    public Boolean updateTaskStatus(AwkIndexTask awkIndexTask) {
        return this.updateById(awkIndexTask);
    }
}
